package dst.search_tree.avl;

import dst.search_tree.avl.BinarySearchTree.TreeNode;

/**
 *
 * Fully Tested: No
 * Fully Documented: No
 * Created: 19-Mar-2012
 * Last Updated: 19-Mar-2012
 *
 * @param <T> 
 * @author Shimu Wu
 */
public class AvlTree<T extends Comparable<T>> extends BinarySearchTree<T> {

    /**
     * A sentinel TreeNode that servers as null, NIL, or the children of a 
     * leaf node
     */
    private final AvlNode SENTINEL = new AvlNode(null);
    
    private AvlNode root;

    /**
     * Initialize an AvlTree with the root containing the given data
     * 
     * @param rootData 
     */
    public AvlTree(T rootData) {
        this.root = new AvlNode(rootData);
        this.size = 1;
    }

    /**
     *
     * Fully Tested: No
     * Fully Documented: No
     * Created: 22-Mar-2012
     * Last Updated: 22-Mar-2012
     *
     * @param <T> 
     * @author Shimu Wu
     */
    protected class AvlNode extends TreeNode {

        private int balanceFactor;

        /**
         * 
         * @param data
         */
        private AvlNode(T data) {
            super(data);
            this.balanceFactor = 0;
        }

        @Override
        protected AvlNode left() {
            return (AvlNode) super.left();
        }

        @Override
        protected AvlNode right() {
            return (AvlNode) super.right();
        }

        /**
         * Rotate the tree rooted at this node to the right.
         * 
         * Precondition: called on a node that is left heavy 
         * (performBalance factor of -1 or -2) and its left child is also left
         * heavy or balanced (performBalance factor of 0 or -1)
         */
        private void rotateRight(AvlNode parent) {
            // There is no checking for SENTINEL since you'll never call
            // a right rotation if the left child is SENTINEL
            /*
             * Before:
             *                       P (parent)
             *                    /      \
             *                   /        \ 
             *                  /          \
             *                 A (this)    /\
             *               /    \       /T0\
             *              /      \     ------
             * (this.left) B       /\
             *           /   \    /T1\
             *          /     \  ------
             *         /\     /\
             *        /T2\   /T3\
             *       ------ ------
             * 
             * After:
             *                       P (parent)
             *                    /      \
             *                   /        \ 
             *                  /          \
             *     (this.left) B           /\
             *              /     \       /T0\
             *             /       \     ------
             *            /\        A (this)
             *           /T2\     /   \   
             *          ------   /     \
             *                  /\     /\
             *                 /T3\   /T1\
             *                ------ ------
             *              
             */
            AvlNode A, B;

            A = this;
            B = this.left();

            if (parent.exists()) {
                if (parent.left() == A) {
                    parent.setLeft(B);
                } else {
                    parent.setRight(B);
                }
            }

            A.setLeft(B.right());
            B.setRight(A);
        }

        private void balanceRightHeavy(AvlNode parent) {
            AvlNode A, B;
            A = this;
            B = A.right();


            if (B.balanceFactor == 1) { // right heavy
                /*
                 *            A (this, +2)
                 *          /    \
                 *         /      \
                 *        /\       B (this.right, +1)
                 *       /T1\    /   \
                 *      ------  /     \
                 *             /\     /\
                 *            /T2\   /T3\
                 *           ------ /    \
                 *                 --------
                 * 
                 */
                A.rotateLeft(parent);
                A.balanceFactor = 0;
                B.balanceFactor = 0;

                if (A == root) {
                    root = B;
                }
            } else { // left heavy
                /*
                 *            A (this, +2)
                 *          /    \
                 *         /      \
                 *        /\       \     
                 *       /T1\        B (this.right, -1)
                 *      ------     /   \
                 *                /     \
                 *               C      /\
                 *             /   \   /T3\ 
                 *            /     \ ------
                 *          / \     / \
                 *         /T21\   /T22\
                 *        ------- ------- 
                 *                 
                 */


                AvlNode C = B.left();

                B.rotateRight(A);
                A.rotateLeft(parent);

                if (C.balanceFactor == 0) {
                    A.balanceFactor = 0;
                    B.balanceFactor = 0;

                } else if (C.balanceFactor == 1) {
                    A.balanceFactor = -1;
                    B.balanceFactor = 0;

                } else {
                    A.balanceFactor = 0;
                    B.balanceFactor = 1;
                }
                C.balanceFactor = 0;

                if (A == root) {
                    root = C;
                }
            }

        }

        private void balanceLeftHeavy(AvlNode parent) {
            AvlNode A, B;
            A = this;
            B = A.left();

            if (B.balanceFactor == -1) { // left heavy
                A.rotateRight(parent);
                A.balanceFactor = 0;
                B.balanceFactor = 0;

                if (A == root) {
                    root = B;
                }
            } else {
                AvlNode C = B.right();

                B.rotateLeft(A);
                A.rotateRight(parent);

                if (C.balanceFactor == 0) {
                    A.balanceFactor = 0;
                    B.balanceFactor = 0;

                } else if (C.balanceFactor == 1) {
                    A.balanceFactor = 0;
                    B.balanceFactor = -1;

                } else {
                    A.balanceFactor = 1;
                    B.balanceFactor = 0;
                }
                C.balanceFactor = 0;

                if (A == root) {
                    root = C;
                }
            }
        }

        /**
         * 
         * @param data
         * @param parent
         * @return the change in height of the tree rooted at this node
         */
        private int add(T data, AvlNode parent) {
            if (this.greaterThan(data)) {
                // current.data > data, go left

                if (!this.hasLeftChild()) {
                    this.setLeft(new AvlNode(data));

                } else if (this.left().add(data, this) == 0) {
                    // If adding the node causes the left subtree to be 
                    // balance, return right away
                    return 0;
                }

                // Otherwise decrement the balance factor (coming from the left)
                this.balanceFactor--;
                if (this.balanceFactor == -2) {
                    this.balanceLeftHeavy(parent);

                    // Once its balanced, the balance factor is 0
                    return 0;

                } else {
                    // this.balanceFactor = 0 or -1
                    // if -1, then the parent needs to update their balance
                    // factor accordingly
                    return Math.abs(this.balanceFactor);
                }

            } else {
                // this.data <= data, go right

                if (!this.hasRightChild()) {
                    this.setRight(new AvlNode(data));

                } else if (this.right().add(data, this) == 0) {
                    // If adding the node causes the right subtree to be 
                    // balance, return right away

                    return 0;
                }

                // Otherwise increment the balance factor (coming from the left)
                this.balanceFactor++;

                if (this.balanceFactor == 2) {
                    this.balanceRightHeavy(parent);
                    return 0;

                } else {
                    // could only be 0 or +1
                    return Math.abs(this.balanceFactor);
                }
            }
        }

        private String getName(AvlNode node) {
            if (node == SENTINEL) {
                return null;
            } else {
                return node.getData().toString();
            }
        }

        @Override
        public String toString() {
            StringBuilder b = new StringBuilder(50);
            b.append("AvlNode: ");
            b.append(this.getData());
            b.append(" balance factor: ");
            b.append(this.balanceFactor);
            b.append(" left: ");
            b.append(getName(this.left()));
            b.append(" right: ");
            b.append(getName(this.right()));

            return b.toString();
        }
    }

    @Override
    protected AvlNode getRoot() {
        return root;
    }

    @Override
    protected AvlNode getSentinel() {
        return SENTINEL;
    }
    
    @Override
    public void add(T data) {
        root.add(data, SENTINEL);
        size++;
    }

    /**
     * Returns the String representation of the root node.
     * 
     * @return 
     */
    @Override
    public String toString() {
        return root.toString();
    }
}